import warnings

warnings.filterwarnings('ignore')

import yaml
from ultralytics import YOLOv10
from pathlib import Path

ROOT = Path(__file__).parent


# 🔥 加载 YAML 文件为字典
def load_hyperparameters(yaml_path: str) -> dict:
    with open(yaml_path, 'r', encoding='utf-8') as f:
        hyp = yaml.safe_load(f)
    print("✅ 成功加载超参数配置:", hyp)
    return hyp


if __name__ == '__main__':
    # 加载超参数字典
    hyp_path = ROOT / 'hyp.VOC.yaml'
    hyperparameters = load_hyperparameters(str(hyp_path))

    # 初始化 YOLOv10 模型
    model = YOLOv10(ROOT / 'ultralytics/cfg/models/v10/yolov10m.yaml')

    # 加载预训练的权重
    model.load(ROOT / 'yolov10m.pt')

    # 训练模型时应用超参数字典
    model.train(
        data=ROOT / 'ultralytics/cfg/datasets/cell.yaml',  # 数据集配置文件
        epochs=200,
        batch=16,
        imgsz=640,
        optimizer='AdamW',
        patience=50,
        device='0',
        amp=False,  # 禁用 AMP (自动混合精度)

        # 通过解包字典的方式传递超参数
        **hyperparameters
    )
